from Transformer_decoder import *

# 解码器类 Generator 实现思路分析
# init函数 (self, d_model, vocab_size)
    # 定义线性层self.project
# forward函数 (self, x)
    # 数据 F.log_softmax(self.project(x), dim=-1)

class Generator(nn.Module):
    def __init__(self,embedding_dim,vocabulary_size):
        # 参数d_model 线性层输入特征尺寸大小
        # 参数vocab_size 线层输出尺寸大小

        super().__init__()
        # 定义线性层
        self.project = nn.Linear(embedding_dim,vocabulary_size)

    def forward(self,y):
        # 数据经过线性层 最后一个维度归一化 log方式
        y = F.log_softmax(self.project(y),dim=-1)
        return y

def test_output():
    decoder_result = test_decoder()
    generator = Generator(embedding_dim,vocabulary_size)

    generator_result = generator(decoder_result)
    print(f'输出后的形状{generator_result.shape}')

if __name__ == '__main__':
    test_output()